#!/usr/bin/env perl

# Copyright (C) Yichun Zhang (agentzh)
# Copyright (C) Guanlan Dai

use 5.006001;
use strict;
use warnings;

use Getopt::Std qw( getopts );

my %opts;

getopts("a:dhp:", \%opts)
    or die usage();

if ($opts{h}) {
    print usage();
    exit;
}

my $pid = $opts{p}
    or die "No nginx process pid specified by the -p option\n";

if ($pid !~ /^\d+$/) {
    die "Bad -p option value \"$pid\": not look like a pid\n";
}

my $stap_args = $opts{a} || '';

if ($^O ne 'linux') {
    die "Only linux is supported but I am on $^O.\n";
}

my $exec_file = "/proc/$pid/exe";
if (!-f $exec_file) {
    die "Nginx process $pid is not running or ",
        "you do not have enough permissions.\n";
}

my $nginx_path = readlink $exec_file;

my $ver = `stap --version 2>&1`;
if (!defined $ver) {
    die "Systemtap not installed or its \"stap\" utility is not visible to the PATH environment: $!\n";
}

if ($ver =~ /version\s+(\d+\.\d+)/i) {
    my $v = $1;
    if ($v < 2.1) {
        die "ERROR: at least systemtap 2.1 is required but found $v\n";
    }

} else {
    die "ERROR: unknown version of systemtap:\n$ver\n";
}

my $stap_src;

my $cmcf = qq{\@cast(cmcf, "ngx_http_core_main_conf_t", "$nginx_path")};
my $ph = qq{\@cast(ph, "ngx_http_phase_handler_t", "$nginx_path")};
my $handler_addr = qq{\@cast(handler_addr, "ngx_http_handler_pt", "$nginx_path")};
my $sizeof_ngx_http_phase_handler_t = qq{&\@cast(0,
                                                "ngx_http_phase_handler_t",
                                                "$nginx_path")[1]};
my $sizeof_ngx_http_handler_pt = qq{&\@cast(0, "ngx_int_t", "$nginx_path")[1]};

my $NGX_HTTP_POST_READ_PHASE = 0;
my $NGX_HTTP_SERVER_REWRITE_PHASE = 1;
my $NGX_HTTP_FIND_CONFIG_PHASE = 2;
my $NGX_HTTP_REWRITE_PHASE = 3;
my $NGX_HTTP_POST_REWRITE_PHASE = 4;
my $NGX_HTTP_PREACCESS_PHASE = 5;
my $NGX_HTTP_ACCESS_PHASE = 6;
my $NGX_HTTP_POST_ACCESS_PHASE = 7;
my $NGX_HTTP_TRY_FILES_PHASE = 8;
my $NGX_HTTP_CONTENT_PHASE = 9;
my $NGX_HTTP_LOG_PHASE = 10;

my $preamble = <<_EOC_;
probe begin {
    printf("Tracing %d ($nginx_path)...\\n", target())
}
_EOC_
chop $preamble;

$stap_src = <<_EOC_;
$preamble

function print_phase(cmcf, phase_num) {
    n = $cmcf->phases[phase_num]->handlers->nelts
    handler_addr = $cmcf->phases[phase_num]->handlers->elts
    handler = $handler_addr
    for (i = 0; i < n; i++) {
        printf("    %s\\n", usymname(handler))
        handler_addr += $sizeof_ngx_http_handler_pt
        handler = $handler_addr
    }
}

probe process("$nginx_path").function("ngx_http_log_request")
{
    if (pid() == target()) {
        begin = gettimeofday_us()
        ngx_http_core_module = &\@var("ngx_http_core_module")
        ctx_index = \@cast(ngx_http_core_module, "ngx_module_t")->ctx_index
        cmcf = \$r->main_conf[ctx_index]
        use_rewrite = $cmcf->phases[$NGX_HTTP_REWRITE_PHASE]->handlers->nelts ? 1 : 0
        use_access = $cmcf->phases[$NGX_HTTP_ACCESS_PHASE]->handlers->nelts ? 1 : 0
        n = 0

        for (i = 0; i < $NGX_HTTP_LOG_PHASE; i++) {
            if (i == $NGX_HTTP_FIND_CONFIG_PHASE) {
                n++
                continue
            }

            if (i == $NGX_HTTP_POST_REWRITE_PHASE) {
                if (use_rewrite) {
                    n++
                }
                continue
            }

            if (i == $NGX_HTTP_POST_ACCESS_PHASE) {
                if (use_access) {
                    n++
                }
                continue
            }

            if (i == $NGX_HTTP_TRY_FILES_PHASE) {
                if ($cmcf->try_files) {
                    n++
                }
                continue
            }

            phase_len = $cmcf->phases[i]->handlers->nelts
            if (phase_len == 0) {
                continue
            }

            if (i == $NGX_HTTP_POST_READ_PHASE) {
                println("post-read phase")

            } else if (i == $NGX_HTTP_SERVER_REWRITE_PHASE) {
                println("server-rewrite phase")

            } else if (i == $NGX_HTTP_REWRITE_PHASE) {
                println("rewrite phase")

            } else if (i == $NGX_HTTP_PREACCESS_PHASE) {
                println("pre-access phase")

            } else if (i == $NGX_HTTP_ACCESS_PHASE) {
                println("access phase")

            } else if (i == $NGX_HTTP_CONTENT_PHASE) {
                println("content phase")

                if (\$r->content_handler) {
                    printf("    %s (request content handler)\\n",
                           usymname(\$r->content_handler))
                }

            } else {
                printf("unknown phase (%d)\\n", i)
            }

            for (j = 0; j < phase_len; j++) {
                handler = $cmcf->phase_engine->handlers[n++]->handler
                if (handler == 0) {
                    continue
                }
                printf("    %s\\n", usymname(handler))
            }

            printf("\\n")
        }

        println("log phase")
        print_phase(cmcf, $NGX_HTTP_LOG_PHASE)

        elapsed = gettimeofday_us() - begin
        printf("\\n%d microseconds elapsed in the probe handler.\\n", elapsed)
        exit()
    }
}
_EOC_

if ($opts{d}) {
    print $stap_src;
    exit;
}

open my $in, "|stap --skip-badvars $stap_args -x $pid -"
    or die "Cannot run stap: $!\n";

print $in $stap_src;

close $in;

sub usage {
    return <<'_EOC_';
Usage:
    ngx-phase-handlers [optoins]

Options:
    -a <args>           Pass extra arguments to the stap utility.
    -d                  Dump out the systemtap script source.
    -h                  Print this usage.
    -p <pid>            Specify the nginx worker process pid.

Examples:
    ngx-phase-handlers -p 12345
_EOC_
}

